from abc import ABCMeta, abstractmethod
import math
from scipy.optimize import leastsq, curve_fit
import numpy as np


class LeastSquares(object):
    def __init__(self):
        self.attributes = "LeastSquares"

    # @abstractmethod
    def print_attributes(self):
        print(self.attributes)

    @abstractmethod
    def fitting(self, raw_data):
        pass

    def func(self, p, x):
        pass

    def error(self, p, x, y):
        return self.func(p, x) - y

    def res_calculate(self, raw_data, p):
        x = raw_data[:, 0]
        y = raw_data[:, 1]
        self.xy_res = []
        # 完全平方
        for d in range(0, len(x)):
            self.xy_res.append(self.error(p[0], x[d], y[d]))

        # 计算残差平方和，和越小表明拟合的情况越好
        self.xy_res_square_sum = np.dot(self.xy_res, self.xy_res) / len(x)
        return self.xy_res_square_sum

    def points_re_judgement(self, outliers_points, p):
        right_points = []
        error_points = []
        for point in outliers_points:
            res = self.error(p[0], point[0], point[1])
            if res * res <= 2 * max(self.xy_res) * max(self.xy_res):
                right_points.append(point)
            else:
                error_points.append(point)
        return right_points, error_points

    def complete_curve(self, all_data, p):
        all_fitted = self.func(p[0], all_data[:, 0])  # 画出拟合后的曲线
        return all_fitted


class LinerPolynomial(LeastSquares):
    def __init__(self):
        super(LinerPolynomial, self).__init__()
        self.attributes = "LeastSquares"

    def print_attributes(self):
        print(self.attributes)

    def fitting(self, raw_data):
        p0 = np.array([1, 20])
        x_raw = raw_data[:, 0]
        y_raw = raw_data[:, 1]

        parameter = leastsq(self.error, p0, args=(x_raw, y_raw))
        y_fitted = self.func(parameter[0], x_raw)  # 画出拟合后的曲线
        return y_fitted, parameter

    def func(self, p, x):
        a1, a2 = p
        return a1 * x + a2


class BinomialPolynomial(LeastSquares):
    def __init__(self):
        super(BinomialPolynomial, self).__init__()
        self.attributes = "BinomialPolynomial"

    def print_attributes(self):
        print(self.attributes)

    def fitting(self, raw_data):
        p0 = np.array([0.1, -0.01, 100])
        x_raw = raw_data[:, 0]
        y_raw = raw_data[:, 1]

        parameter = leastsq(self.error, p0, args=(x_raw, y_raw))
        y_fitted = self.func(parameter[0], x_raw)  # 画出拟合后的曲线
        return y_fitted, parameter

    def func(self, p, x):
        a1, a2, a3 = p
        return a1 * x**2 + a2 * x + a3


class TriPolynomial(LeastSquares):
    def __init__(self):
        super(TriPolynomial, self).__init__()
        self.attributes = "TriPolynomial"

    def print_attributes(self):
        print(self.attributes)

    def fitting(self, raw_data):
        p0 = np.array([0.1, -0.01, 1, 10])
        x_raw = raw_data[:, 0]
        y_raw = raw_data[:, 1]

        parameter = leastsq(self.error, p0, args=(x_raw, y_raw))
        y_fitted = self.func(parameter[0], x_raw)  # 画出拟合后的曲线
        return y_fitted, parameter

    def func(self, p, x):
        a1, a2, a3, a4 = p
        return a1 * x**3 + a2 * x**2 + a3 * x + a4


class CurveFit:
    def __init__(self, raw_arr):
        self.raw_array = raw_arr
        self.ploy_factory = LeastSquares()

    def fit_process(self, polynomial):
        if polynomial == 1:
            self.ploy_factory = LinerPolynomial()
        elif polynomial == 2:
            self.ploy_factory = BinomialPolynomial()
        elif polynomial == 3:
            self.ploy_factory = TriPolynomial()

        # ploy_factory.print_attributes()
        self.y_fitted, self.parameter = self.ploy_factory.fitting(self.raw_array)
        return self.y_fitted

    def calculate_residual(self):
        sum_res = self.ploy_factory.res_calculate(self.raw_array, self.parameter)
        return sum_res

    def outliers_re_judgment(self, outliers_points):
        right_points, error_points = self.ploy_factory.points_re_judgement(outliers_points, self.parameter)
        return right_points, error_points

    def complete_curve(self, all_data):
        all_curve = self.ploy_factory.complete_curve(all_data, self.parameter)
        return all_curve
